# backend/asr.py
import logging
import numpy as np
from funasr import AutoModel

logger = logging.getLogger("asr")

class StreamingASR:
    def __init__(self, model_name="paraformer-zh-streaming", revision="v2.0.4"):
        logger.info(f"Loading ASR model {model_name} ...")
        # AutoModel may load heavy model — ensure environment has resources
        self.model = AutoModel(model=model_name, model_revision=revision)
        self.sessions = {}

    def start(self, session_id: str):
        self.sessions[session_id] = {}

    def accept_chunk(self, session_id: str, pcm_int16: np.ndarray, is_final: bool = False) -> str:
        cache = self.sessions.get(session_id, {})
        try:
            res = self.model.generate(pcm_int16, cache=cache, is_final=is_final)
            if res and len(res) > 0:
                return res[0].get("text", "") or ""
        except Exception:
            logger.exception("ASR generate error")
        return ""

    def finalize(self, session_id: str) -> str:
        cache = self.sessions.get(session_id, {})
        try:
            res = self.model.generate(np.array([], dtype=np.int16), cache=cache, is_final=True)
            if res and len(res) > 0:
                text = res[0].get("text", "") or ""
            else:
                text = ""
        except Exception:
            logger.exception("ASR finalize error")
            text = ""
        # cleanup session
        self.sessions.pop(session_id, None)
        return text
